"""Contains a class, VpfitModel, useful for parsing f26 and fort.13
files, and writing out fort.13 files.
"""
import numpy as np

Ckms = 299792.458   # speed light km/s. exact

# the data types of the lines and regions numpy arrays:

dtype_lines = [('name', 'S6'),
               ('z', 'f8'),
               ('zpar', 'S2'),
               ('b', 'f8'),
               ('bpar', 'S2'),
               ('logn', 'f8'),
               ('lognpar', 'S2'),
               ('zerr', 'f8'),
               ('berr', 'f8'),
               ('lognerr', 'f8')]

dtype_regions = [('filename', 'S150'),
                 ('num', 'S1'),
                 ('wmin', 'f8'),
                 ('wmax', 'f8'),
                 ('resolution', 'S11')]

def _parselines(params):
    """ Separates the parameters from their tied/fixed/special
    characters.
    """
    temp = []
    for name,redshift,b,logn,zerr,berr,lognerr in params:
        i = -1
        while not redshift[i].isdigit():  i -= 1
        if i != -1:
            redshift,redshift_par = float(redshift[:i+1]),redshift[i+1:]
        else:
            redshift,redshift_par = float(redshift),''
        i = -1
        while not b[i].isdigit():  i -= 1
        if i != -1:
            b,b_par = float(b[:i+1]), b[i+1:]
        else:
            b,b_par = float(b),''
        i = -1
        while not logn[i].isdigit():  i -= 1
        if i != -1:
            logn,logn_par = float(logn[:i+1]),logn[i+1:]
        else:
            logn,logn_par = float(logn),''
        try:
            zerr = float(zerr)
        except ValueError:
            zerr = -1
        try:
            berr = float(berr)
        except ValueError:
            berr = -1
        try:
            lognerr = float(lognerr)
        except ValueError:
            lognerr = -1

        temp.append((name,redshift,redshift_par,b,b_par,logn,logn_par,
                     zerr,berr,lognerr))

    foo = np.array(temp, dtype=dtype_lines)
    return foo.view(np.rec.recarray)

def _parseregions(rows, res=None):
    """ Parses the region information from a f26 or fort.13 file. """
    if res is None:
        res = ''
    rinfo = []
    for row in rows:
        r = row.split('!')[0].lstrip().lstrip('%%').split()
        nitems = len(r)
        r[2] = float(r[2])
        r[3] = float(r[3])
        if nitems == 4:
            rinfo.append(tuple(r + [res]))
        elif nitems == 5:
            rinfo.append(tuple(r))
        else:
            raise Exception('bad format in fitting regions:\n %s' % row)

    foo = np.array(rinfo, dtype=dtype_regions)
    return foo.view(np.rec.recarray)


def sumlines(lines):
    """ Given several lines (record array), returns them in the vpfit
    summed format. """

    summedlines = lines.copy()

    logntots = np.log10(np.sum(10**lines.logn))

    for i,logntot in enumerate(logntots):
        if i == 0:
            summedlines[i].logn = logntot
            #summedlines[i].lognstr = '%7.4f' % logntot
        summedlines[i].lognpar = 'w'
    
    return summedlines

class VpfitModel(object):
    """
    All the info about a vpfit model. Can read fort.13 and f26 files
    (old and new style).  Can write out the model as a fort.13 style file.
    """
    def __init__(self, names=None, logn=None, z=None, b=None,
                 zpar=None, bpar=None, lognpar=None,
                 filenames=None, wmin=None, wmax=None, res=None, num=None):
        if None in (names,logn,z,b):
            self.lines = None              # record array
        else:
            ncomp = len(z)
            if zpar is None:  zpar = [''] * ncomp
            if bpar is None:  bpar = [''] * ncomp
            if lognpar is None:  lognpar = [''] * ncomp
            zerr = [-1] * ncomp
            berr = [-1] * ncomp
            lognerr = [-1] * ncomp
            foo = np.rec.fromarrays([names,z,zpar,b,bpar,logn,lognpar,zerr,
                                     berr,lognerr], dtype=dtype_lines)
            self.lines = foo
        if None in (filenames,wmin,wmax):
            self.regions = None            # record array
        else:
            if res is None:
                res = [''] * len(filenames)
            if num is None:
                num = ['1'] * len(filenames)
            foo = np.rec.fromarrays([filenames,num,wmin,wmax,res],
                                    dtype=dtype_regions)
            self.regions = foo
        self.stats = None
        self.iden = None        # A unique id for this model


    def readf26(self, filename, res=None):
        """ Reads in a f26 style file. if the keyword res is given,
        this string provides the resolution information for the
        spectra fitted. For example: res='vsig=69.0' """
        f = open(filename).readlines()
        regionrows = [r for r in f if r.lstrip().startswith('%%')]
        ionrows = [r for r in f if '%%' not in r and 'Stats' not in r]
        statrow = [row for row in f if 'Stats' in row]
        if statrow:
            status = 'BAD' if statrow[0].split()[-1] == 'BAD' else 'OK'
            vals = statrow[0].split()[2:] + [status]
        else:
            # older style f26 file
            stat = ionrows[0]
            status = 'BAD' if stat.split()[-1] == 'BAD' else 'OK'
            vals = [stat[66:71], stat[71:85], stat[85:90], stat[90:95],
                    stat[95:102], stat[102:107], status]
        keys = 'iterations nchisq npts dof prob ndropped info'.split()
        self.stats = dict(zip(keys,vals))
        self.regions = _parseregions(regionrows,res=res)
        #print self.regions,'\n\n\n'
        ionrows = [r.lstrip() for r in ionrows]
        param = [ [r[:5]] + r[5:].split() for r in ionrows]
        param = [[p[0],p[1],p[3],p[5],p[2],p[4],p[6]] for p in param]
        self.lines = _parselines(param)

    def readf13(self, filename, read_regions=True, res=None):
        """ Reads in a fort.13 style file. """
        f = open(filename).readlines()
        f = [row.lstrip() for row in f[1:]]      # skip past first line with *
        isep = [row[0] for row in f].index('*')  # find separating *
        if read_regions:
            self.regions = _parseregions([row for row in f[:isep]],res=res)
        param  = [[row[:5]] + row[5:].split() for row in f[isep+1:]]
        param = [[p[0],p[2],p[3],p[1],-1,-1,-1] for p in param]
        self.lines = _parselines(param)
        self.stats = None

    def writef26(self,filename, write_regions=True):
        """ Writes out a f26 style file."""
        temp = []
        if write_regions and self.regions is not None:
            for r in self.regions:
                temp.append('%%%% %(filename)s  %(num)s  %(wmin)7.2f '
                            '%(wmax)7.2f  %(resolution)s\n' % r)
        if self.lines is not None:
            for line in self.lines:
                temp.append('   %(name)s     %(z)11.8f%(zpar)-2s '
                            '%(zerr)11.8f %(b)6.2f%(bpar)-2s %(berr)6.2f '
                            '%(logn)7.4f%(lognpar)-2s %(lognerr)7.4f\n' % line)
        open(filename,'w').writelines(temp)

    def writef13(self, filename, write_regions=True):
        """ Writes out a fort.13 style file. """
        # The whitespace is important if the f13 files are to be read
        # by vpguess - don't change it!
        temp = []
        if write_regions:
            temp.append('   *\n')
            if self.regions is not None:
                for r in self.regions:
                    temp.append('%(filename)s  %(num)s  %(wmin)7.2f '
                                '%(wmax)7.2f  %(resolution)s\n' % r)
            temp.append('  *\n')
        if self.lines is not None:
            for line in self.lines:
                temp.append('   %(name)s     %(logn)7.4f%(lognpar)-2s  '
                            '%(z)11.8f%(zpar)-2s  %(b)6.2f%(bpar)-2s '
                            '0.00   0.00E+00  0\n' % line)
        open(filename,'w').writelines(temp)


def read_26info(resultdir):
    """Read in a series of f26 result files.

    All those that are from 'good' fits, lump together.
    """
    prefix = os.path.abspath(resultdir) + '/'
    filenames = [n for n in os.listdir(resultdir) if n.endswith('f26')]
    f26 = []
    lines = []
    rejected = 0

    for filename in sorted(filenames):
        vp = VpfitModel()
        vp.readf26(prefix + filename)
        vp.iden = prefix + filename
        f26.append(vp)
        # if number of lines is suspiciously low, or the chisq didn't
        # reach the threshold, don't include lines in total list.
        #if ((len(vp.lines) < 4) or
        if float(vp.stats['nchisq']) > 1.1:
            print len(vp.lines), float(vp.stats['nchisq']), filename
            rejected += 1
            continue
        cond = vp.lines.name != '<>'
        l = vp.lines[cond]
        lines.extend(l.tolist())

    print '%d systems rejected' % rejected
    lines = np.array(lines, dtype=f26[0].lines.dtype)

    return f26, lines.view(np.rec.recarray)

def calc_v90(vp, plot=False, z0=None,
             wav0=1215.6701, osc=0.4164, gam=6.265e8):
    """ For a vp model, we want to calculate the velocity width that
    contains 90% of the the total optical depth at the lya line (or
    perhaps it is the same regardless of which transition I take?) v_90
    is defined in Prochaska and Wolfe 1997.

    At the moment it guesses how big a velocity range it has to
    calculate the optical depth over - a bit dodgy"""
    lines = vp.lines
    #print 'calculating for %s' % lines
    # work in velocity space
    z = lines.z
    if z0 is None:  z0 = np.median(z)
    vel = (z - z0) / (1 + z0) * Ckms
    # calculate the optical depth as a function of velocity, 500 km/s
    # past the redmost and bluemost components - hopefully this is far
    # enough (maybe not for DLAs?)
    dv = 0.5
    vhalf = (vel.max() - vel.min())/2. + 300
    v = np.arange(-vhalf, vhalf + dv, dv)
    tau = np.zeros(len(v))
    for line,vline in zip(lines,vel):
        if line['logn'] > 21.0:
            print ('very (too?) high logN: %s' % line['logn'])
            print ('returning width of -1')
            return -1.
        temptau = calctau(v, vline, wav0, osc, gam, line['logn'],
                          btemp=line['b'])
        tau += temptau
        #pl.plot(v,tau,'+-')
        #raw_input('N %(logn)f b %(b)f enter to continue' % line)

    # integrate over the entire v range to calculate integral of tau wrt v.
    sumtaudv = np.trapz(tau,dx=dv)
    lenv = len(v)
    # starting from the left v edge, increase v until int from left
    # edge to v gives 5% of total integral
    sum5perc = sumtaudv / 20.
    sumtau = 0.
    i = 0
    while (sumtau < sum5perc):
        i += 1
        sumtau = np.trapz(tau[:i])
        if i == lenv:
            raise Exception('Problem with velocity limits!')
    vmin = v[i-1]
    # Do the same starting from the right edge.
    sumtau = 0
    i = -1
    while (sumtau < sum5perc):
        sumtau = np.trapz(tau[i:])
        i -= 1
        if -i == lenv:
            raise Exception('Problem with velocity limits!')
    vmax = v[i+1]
    # Difference between the two is v_90
    v90 = vmax - vmin
    if plot:
        pl.plot(v,tau,'+-')
        pl.vlines((vmin,vmax),0,tau.max())
    #raw_input('Enter to continue...')
    return v90



# if __name__ == '__main__':
#     m = VpfitModel()

#     m.readf26(sys.argv[1])

#     for i,line in enumerate(m.lines):
#         if 'y' not in line.logn:  continue
#         hcolsum = float(line.logn.replace('y',''))
#         hcol2 = float(lines.logn[i+1].replace('y',''))
#         lines.logn[i] = '%9.6f' % (log10(10**hcolsum - 10**hcol2))
#         lines.logn[i+1] = lines.logn[i+1].replace('y','')
#         dcolsum = float(lines.logn[i+2].replace('%',''))
#         dcol2 = float(lines.logn[i+3].replace('Y',''))
#         lines.logn[i+2] = '%9.6f' % (log10(10**dcolsum - 10**dcol2))
#         lines.logn[i+3] = lines.logn[i+3].replace('Y','')
#         break

#     # convert d/h lines with summed column densities to normal column
#     # densities so vpguess can read the fort.13 file

#     # Now write out fort.13

#     m.writef13(sys.argv[2])
